#!/usr/bin/env python

"""
pre-commit hook to check repo tree

To use it normally as git pre-commit hook,
make a symlink and install dependencies:
$ ln -s ../../pre-commit .git/hooks
$ pacman -Syu --needed python-yaml python-jsonschema pyalpm

To check modified packages, run without options:
$ ./pre-commit

To fully test all packages, run with "-a" option:
$ ./pre-commit -a
"""

import os
import sys
import ast
import gzip
import json
from time import strftime
from subprocess import Popen, PIPE
from unittest import TestCase, TestLoader, TextTestRunner, skip
from io import StringIO
from typing import Generator

# need python-yaml package
import yaml

# need python-jsonschema package
import jsonschema

DIRS, FILES, COMMITS = None, None, None
GIT_BRANCH = 'HEAD'
ORIGIN_BRANCH = 'origin/master'
CHANGED_PACKAGES = set()
SUBFOLDER = 'aur-repo'

# URL for downloading latest database of packages and groups in official repos
DUMP_GROUPS_URL = 'https://build.archlinuxcn.org/~farseerfc/dump-groups.gz'
DUMP_GROUPS_TMP = '/tmp/dump-groups-%s.gz'
# Cache downloaded database for 1 hour
DUMP_GROUPS_TIMEFORMAT = '%Y%m%d%H%z'

GROUPS, PACKAGES = {}, {}

YAML_SCHEMA = None

CHECK_ALL = False
IGNORE_SRCINFO = False


def readutf8(st):
    return st.read().decode('utf-8')


def git_write_tree_index():
    global GIT_BRANCH
    cmd = ['git', 'write-tree']
    with Popen(cmd, stdout=PIPE) as p:
        GIT_BRANCH = readutf8(p.stdout).rstrip()


def git_diff_tree():
    global GIT_BRANCH, ORIGIN_BRANCH, CHANGED_PACKAGES
    # git diff-tree
    # -r Recurse into sub-trees.
    # -z \0 line termination on output and do not quote filenames.
    cmd = ['git', 'diff-tree', ORIGIN_BRANCH + '..' + GIT_BRANCH,
           '-r', '-z', '--name-only']
    with Popen(cmd, stdout=PIPE) as p:
        changed_tree = readutf8(p.stdout)
        for line in changed_tree.split('\0')[:-1]:  # ignore last \0
            paths = line.strip().split('/')
            if git_isdir(SUBFOLDER):
                if len(paths) > 2:
                    if git_isdir(os.path.join(SUBFOLDER, paths[1])):
                        CHANGED_PACKAGES.add(paths[1])
            else:
                if len(paths) > 1:
                    if git_isdir(paths[0]):
                        CHANGED_PACKAGES.add(paths[0])


def git_ls_tree():
    cmd = ['git', 'ls-tree', '-rtz', '--full-name', GIT_BRANCH]
    # git ls-tree
    # -r Recurse into sub-trees.
    # -t Show tree entries even when going to recurse them.
    # -z \0 line termination on output and do not quote filenames.
    dirs = []
    files = []
    commits = []
    with Popen(cmd, stdout=PIPE) as p:
        alllines = readutf8(p.stdout)
        for line in alllines.split('\0')[:-1]:  # ignore last \0
            # OUTPUT FORMAT
            # <mode> SP <type> SP <object> TAB <file>
            first, _, path = line.partition('\t')
            _, t, _ = first.split(' ')
            if t == 'tree':
                dirs.append(path)
            elif t == 'blob':
                files.append(path)
            elif t == 'commit':
                commits.append(path)
        return dirs, files, commits


def git_open(filename, decode_errors='strict'):
    cmd = ['git', 'cat-file', 'blob', f'{GIT_BRANCH}:{filename}']
    with Popen(cmd, stdout=PIPE) as p:
        return StringIO(p.stdout.read().decode('utf-8', errors=decode_errors))


def git_listdir(path):
    cmd = ['git', 'ls-tree', '--name-only',
           (f'{GIT_BRANCH}:{path}' if path != '.' else GIT_BRANCH)]
    with Popen(cmd, stdout=PIPE) as p:
        return readutf8(p.stdout).split("\n")[:-1]


def git_isdir(path):
    global DIRS
    return path in DIRS


def git_isfile(path, name):
    global FILES
    return os.path.join(path, name) in FILES


def pkgpath(pkgname):
    if git_isdir(SUBFOLDER):
        return os.path.join(SUBFOLDER, pkgname)
    else:
        return pkgname


def lilac_py_has_field(pkgname, field, testcase):
    if not git_isfile(pkgpath(pkgname), 'lilac.py'):
        return False
    lilac_path = os.path.join(pkgpath(pkgname), 'lilac.py')
    with git_open(lilac_path) as lilac_py:
        try:
            lilac_ast = ast.parse(lilac_py.read())
            # For simplicity we only check top level assignments
            # to find field such as `update_on`, which shoud be enough
            for block in lilac_ast.body:
                if isinstance(block, ast.Assign):
                    for target in block.targets:
                        if target.id == field:
                            return True
        except:
            testcase.fail(msg="{} can not parse".format(lilac_path))
    return False

def lilac_py_has_field_value(pkgname, field, value, testcase):
    if not git_isfile(pkgpath(pkgname), 'lilac.py'):
        return False
    lilac_path = os.path.join(pkgpath(pkgname), 'lilac.py')
    with git_open(lilac_path) as lilac_py:
        try:
            lilac_ast = ast.parse(lilac_py.read())
            # For simplicity we only check top level assignments
            # to find field such as `update_on`, which shoud be enough
            for block in lilac_ast.body:
                if isinstance(block, ast.Assign):
                    for target in block.targets:
                        if target.id == field:
                            return target.value == value
        except:
            testcase.fail(msg="{} can not parse".format(lilac_path))
    return False

def lilac_py_has_update_on(pkgname, testcase):
    return lilac_py_has_field(pkgname, 'update_on', testcase)

def lilac_py_has_managed_false(pkgname, testcase):
    return lilac_py_has_field_value(pkgname, 'managed', 'False', testcase)

def lilac_yaml_has_field(pkgname, field, testcase):
    if not git_isfile(pkgpath(pkgname), 'lilac.yaml'):
        return False
    lilac_path = os.path.join(pkgpath(pkgname), 'lilac.yaml')
    with git_open(lilac_path) as lilac_yaml:
        try:
            lilac_ast = yaml.load(lilac_yaml.read(), Loader=yaml.SafeLoader)
            if field in lilac_ast:
                return True
        except:
            testcase.fail(msg="{} can not parse".format(lilac_path))
    return False

def lilac_yaml_has_field_value(pkgname, field, value, testcase):
    if not git_isfile(pkgpath(pkgname), 'lilac.yaml'):
        return False
    lilac_path = os.path.join(pkgpath(pkgname), 'lilac.yaml')
    with git_open(lilac_path) as lilac_yaml:
        try:
            lilac_ast = yaml.load(lilac_yaml.read(), Loader=yaml.SafeLoader)
            if field in lilac_ast:
                return lilac_ast[field] == value
        except:
            testcase.fail(msg="{} can not parse".format(lilac_path))
    return False

def lilac_yaml_has_update_on(pkgname, testcase):
    return lilac_yaml_has_field(pkgname, 'update_on', testcase)

def lilac_yaml_has_managed_false(pkgname, testcase):
    return lilac_yaml_has_field_value(pkgname, 'managed', False, testcase)

def extract_srcinfo(pkgbuild):
    dir_path = os.path.dirname(os.path.realpath(__file__))
    with Popen(['bash', dir_path + '/parse-pkgbuild', '/dev/stdin'],
               stdin=PIPE, stdout=PIPE) as p:
        outs,_ = p.communicate(input=pkgbuild.encode('UTF-8'))
        return outs.decode('UTF-8')


def depends_strip_version(depends):
    if '>' in depends:
        return depends[:depends.rfind('>')].strip()
    if '<' in depends:
        return depends[:depends.rfind('<')].strip()
    if depends.count('=') > 1:
        return depends[:depends.rfind('=')].strip()
    return depends.strip()


def dump_groups():
    global GROUPS, PACKAGES
    try:
        import pycman
        # now that we are running on archlinux so dump from pacman's syncdb
        handle = pycman.config.init_with_config('/etc/pacman.conf')
        OFFICIAL = ['core', 'extra']
        syncdbs = [db for db in handle.get_syncdbs() if db.name in OFFICIAL]
        for db in syncdbs:
            for pkg in db.pkgcache:
                name = pkg.name
                if name in PACKAGES:
                    PACKAGES[name].append(db.name)
                else:
                    PACKAGES[name] = [db.name, ]
            for grp in db.grpcache:
                name,_ = grp
                if name in GROUPS:
                    GROUPS[name].append(db.name)
                else:
                    GROUPS[name] = [db.name, ]
    except ImportError:
        # not running on an archlinux/pacman so download from URL
        def refresh_dump_groups():
            dump_groups_path = DUMP_GROUPS_TMP % (strftime(DUMP_GROUPS_TIMEFORMAT))
            if os.path.exists(dump_groups_path):
                return gzip.open(dump_groups_path, 'rt')
            import urllib.request
            urllib.request.urlretrieve(DUMP_GROUPS_URL, dump_groups_path)
            return gzip.open(dump_groups_path, 'rt')

        with refresh_dump_groups() as dump:
            result = json.loads(dump.read())
            if result['format_version'] != 2:
                raise ValueError('dump-groups format mismatch')
            GROUPS = result['groups']
            PACKAGES = result['packages']


def load_lilac_yaml_schema():
    global YAML_SCHEMA
    with open('lilac-yaml-schema.yaml', 'r') as schema:
        YAML_SCHEMA = yaml.load(schema, Loader=yaml.SafeLoader)


def iter_packages(all_pkgs: bool = False) -> Generator[str, None, None]:
    packagesfolder = SUBFOLDER
    if not git_isdir(packagesfolder):
        packagesfolder = '.'
    for package in git_listdir(packagesfolder):
        if not git_isdir(pkgpath(package)):
            continue
        if package[0] == '.':
            continue
        if not(package in CHANGED_PACKAGES or CHECK_ALL or all_pkgs):
            continue
        yield package


class DuplicateKeyError(Exception):
    pass

class UniqueKeyLoader(yaml.SafeLoader):
    def construct_mapping(self, node, deep=False):
        mapping = []
        for key_node, value_node in node.value:
            key = self.construct_object(key_node, deep=deep)
            if key in mapping:
                raise DuplicateKeyError(key)
            mapping.append(key)
        return super().construct_mapping(node, deep)

class RepoTreeTest(TestCase):

    @skip("PKGBUILDs may be added later by lilac")
    def test_pkgbuild_exists(self):
        for package in iter_packages():
            with self.subTest(package=package):
                self.assertTrue(git_isfile(pkgpath(package), "PKGBUILD"),
                                msg=('package "\033[1;31m{0}\033[m" does '
                                     'not have PKGBUILD').format(package))

    @skip("lilac can now check and drop packages in official repos")
    def test_package_in_official(self):
        for package in iter_packages():
            with self.subTest(package=package):
                self.assertFalse(package in PACKAGES,
                                msg=('package "\033[1;31m{0}\033[m" exists in '
                                     'official repo').format(package))


    def test_lilac_yaml_schema(self):
        for package in iter_packages():
            if not git_isfile(pkgpath(package), 'lilac.yaml'):
                continue
            with self.subTest(package=package):
                self.assertTrue(lilac_yaml_has_field(package, 'maintainers', self),
                                msg=('package "\033[1;31m{0}\033[m" does not have '
                                   '"maintainers" in '
                                   '"\033[1;31m{0}/lilac.yaml\033[m" '
                                   'file').format(package))
            with self.subTest(package=package):
                lilac_yaml_path = os.path.join(pkgpath(package), 'lilac.yaml')
                with git_open(lilac_yaml_path) as lilac_yaml_file:
                    try:
                        lilac_yaml = yaml.load(lilac_yaml_file.read(), Loader=UniqueKeyLoader)
                    except Exception:
                        self.fail(msg='package "\033[1;31m{0}\033[m"\'s lilac.yaml'
                                  'contains invalid YAML'.format(package))
                    try:
                        jsonschema.validate(lilac_yaml, YAML_SCHEMA)
                    except jsonschema.exceptions.ValidationError:
                        self.fail(msg='package "\033[1;31m{0}\033[m" contains '
                                  'invalid "\033[1;31mlilac.yaml\033[m"'
                                  ''.format(package))

    def test_lilac_py_has_update_on_or_managed_false(self):
        for package in iter_packages():
            #if not any([git_isfile(pkgpath(package), 'lilac.py'),
            #            git_isfile(pkgpath(package), 'lilac.yaml')]):
            #    continue
            with self.subTest(package=package):
                self.assertTrue(any([
                                     lilac_py_has_update_on(package, self),
                                     lilac_yaml_has_update_on(package, self),
                                     lilac_py_has_managed_false(package, self),
                                     lilac_yaml_has_managed_false(package, self),
                                    ]),
                                msg=('package "\033[1;31m{0}\033[m" does not have '
                                   '"update_on" or "managed: false" in '
                                   '"\033[1;31m{0}/lilac.py\033[m" or '
                                   '"\033[1;31m{0}/lilac.yaml\033[m" '
                                   'files').format(package))

    def test_pkgbuild(self):
        def check_var_srcinfo(var, package, against, line):
            if (var + ' = ') not in line:
                return
            with self.subTest(package=package, line=line.strip()):
                for pkg in against.keys():
                    self.assertNotEqual(var + ' = ' + pkg,
                        depends_strip_version(line),
                        msg=('PKGBUILD of package "\033[1;31m{0}\033[m" '
                                'contains "\033[1;31m{1} = {2}\033[m" in '
                                '"{3}" repo').format(package, var, pkg,
                                                     ','.join(against[pkg])))
        for package in iter_packages():
            if IGNORE_SRCINFO:
                continue
            if not git_isfile(pkgpath(package), 'PKGBUILD'):
                # PKGBUILD may be added later by lilac
                continue
            pkgbuild_path = os.path.join(pkgpath(package), 'PKGBUILD')
            print('checking', pkgbuild_path)
            with git_open(pkgbuild_path, decode_errors='replace') as pkgbuild_file:
                pkgbuild = pkgbuild_file.read()
                if not (('replaces' in pkgbuild) or ('groups' in pkgbuild)):
                    # only parse pkgbuild as srcinfo when necessory
                    continue
                srcinfo = extract_srcinfo(pkgbuild)
                for line in srcinfo.split('\n'):
                    check_var_srcinfo('replaces', package, PACKAGES, line)
                    check_var_srcinfo('groups', package, GROUPS, line)

    def test_repo_depends_valid(self):
        pkgs = list(iter_packages())
        all_pkgs = list(iter_packages(all_pkgs=True))
        for package in pkgs:
            with self.subTest(package=package):
                self.assertFalse(
                    lilac_py_has_field(package, 'repo_depends', self),
                    msg=('package "\033[1;31m{0}\033[m" has repo_depends '
                         'in lilac.py').format(package))
                lilac_yaml_path = os.path.join(pkgpath(package), 'lilac.yaml')
                with git_open(lilac_yaml_path) as lilac_yaml_file:
                    lilac_yaml = yaml.load(lilac_yaml_file.read(), Loader=yaml.SafeLoader)
                    for dep in lilac_yaml.get('repo_depends', []):
                        if isinstance(dep, dict):
                            # When `dep` is a dict, this repo_depends entry
                            # has the form `pkgbase: pkgname`. Here I assume
                            # split packages are successfully packaged. so only
                            # the pkgbase needs to be checked.
                            dep_package = list(dep.keys())[0]
                        else:
                            dep_package = dep
                        self.assertTrue(
                            dep_package in all_pkgs,
                            msg=('package "\033[1;31m{0}\033[m" depends on '
                                 '"\033[1;31m{1}\033[m", which does not exist '
                                 'or is unmanaged').format(package, dep_package))

    def test_no_commit_objects(self):
        # Fails if an embedded git repo (e.g., a submodule) is commited
        self.assertFalse(COMMITS, msg='Unexpected commit objects: {!r}'.format(COMMITS))


def run_test(testcase, msg):
    output = StringIO()
    suite = TestLoader().loadTestsFromTestCase(testcase)
    runner = TextTestRunner(output, verbosity=0)
    results = runner.run(suite)
    if not results.wasSuccessful():
        print(output.getvalue())
        print(msg.format(len(results.failures)))
        sys.exit(1)


def usage():
    print(__doc__)


if __name__ == '__main__':
    if '-h' in sys.argv or '--help' in sys.argv:
        usage()
        sys.exit(0)
    if '-a' in sys.argv or '--all' in sys.argv:
        CHECK_ALL = True
    if '--no-srcinfo' in sys.argv:
        IGNORE_SRCINFO = True

    dump_groups()
    load_lilac_yaml_schema()
    git_write_tree_index()
    DIRS, FILES, COMMITS = git_ls_tree()
    if not CHECK_ALL:
        git_diff_tree()
    run_test(RepoTreeTest,
             msg=('\033[1;31mThere are {0} error(s) inside repo, ' +
                  'blocking git commit. Please fix the errors and commit ' +
                  'again\033[m'))
